package top.simonby.iron.core.util;

import lombok.Getter;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;

import static java.util.concurrent.Executors.newFixedThreadPool;


/**
 * 并发测试工具类
 *
 * @author Simon
 */
public class CountDownLatchUtils {

    private CountDownLatch start;
    private CountDownLatch end;
    @Getter
    private int pollSize = 10;

    public CountDownLatchUtils() {
        this(10);
    }

    public CountDownLatchUtils(int pollSize) {
        this.pollSize = pollSize;
        start = new CountDownLatch(1);
        end = new CountDownLatch(pollSize);
    }

    public void latch(MyFunctionalInterface functionalInterface) throws InterruptedException {
        ExecutorService executorService = newFixedThreadPool(pollSize);
        for (int i = 0; i < pollSize; i++) {
            Runnable run = () -> {
                try {
                    start.await();
                    functionalInterface.run();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                } finally {
                    end.countDown();
                }
            };
            executorService.submit(run);
        }

        start.countDown();
        end.await();
        executorService.shutdown();
    }

    @FunctionalInterface
    public interface MyFunctionalInterface {
        /**
         * 执行
         *
         * @throws InterruptedException
         */
        void run() throws InterruptedException;
    }
}
